"""
情感分类器模块
功能：实现基于传统机器学习的语音情感识别分类器
参数：特征向量、情感标签
返回值：预测的情感类别和置信度
使用场景：语音情感识别的核心分类组件
"""

import joblib
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import classification_report, confusion_matrix
from typing import Tuple, Dict, Any, List
import os

from core.config import settings

class EmotionClassifier:
    """
    语音情感分类器类
    用途：训练和预测语音情感类别
    """
    
    def __init__(self, model_type: str = "random_forest"):
        """
        初始化分类器
        参数：model_type - 模型类型 ("svm", "knn", "random_forest")
        """
        self.model_type = model_type
        self.model = None
        self.scaler = StandardScaler()
        self.emotion_labels = settings.EMOTION_LABELS
        self.is_trained = False
        
        # 根据模型类型初始化分类器
        self._initialize_model()
    
    def _initialize_model(self):
        """
        根据指定类型初始化机器学习模型
        """
        if self.model_type == "svm":
            self.model = SVC(
                kernel='rbf',
                C=1.0,
                gamma='scale',
                probability=True,  # 启用概率预测
                random_state=42
            )
        elif self.model_type == "knn":
            self.model = KNeighborsClassifier(
                n_neighbors=5,
                weights='distance'
            )
        elif self.model_type == "random_forest":
            self.model = RandomForestClassifier(
                n_estimators=100,
                max_depth=10,
                random_state=42
            )
        else:
            raise ValueError(f"不支持的模型类型: {self.model_type}")
    
    def train(self, features: np.ndarray, labels: np.ndarray) -> Dict[str, Any]:
        """
        训练情感分类模型
        参数：
            features - 特征矩阵 (n_samples, n_features)
            labels - 标签数组 (n_samples,)
        返回值：训练结果统计信息
        """
        # 数据预处理
        X_scaled = self.scaler.fit_transform(features)
        
        # 划分训练集和验证集
        X_train, X_val, y_train, y_val = train_test_split(
            X_scaled, labels, 
            test_size=0.2, 
            random_state=42,
            stratify=labels
        )
        
        # 训练模型
        self.model.fit(X_train, y_train)
        self.is_trained = True
        
        # 评估模型
        train_score = self.model.score(X_train, y_train)
        val_score = self.model.score(X_val, y_val)
        
        # 交叉验证
        cv_scores = cross_val_score(self.model, X_scaled, labels, cv=5)
        
        # 预测验证集
        y_pred = self.model.predict(X_val)
        
        return {
            "model_type": self.model_type,
            "train_accuracy": train_score,
            "validation_accuracy": val_score,
            "cv_mean_score": np.mean(cv_scores),
            "cv_std_score": np.std(cv_scores),
            "classification_report": classification_report(y_val, y_pred),
            "confusion_matrix": confusion_matrix(y_val, y_pred).tolist()
        }
    
    def predict(self, features: np.ndarray) -> Dict[str, Any]:
        """
        预测单个样本的情感类别
        参数：features - 特征向量
        返回值：包含预测结果的字典
        """
        if not self.is_trained:
            raise ValueError("模型尚未训练，请先调用train()方法")
        
        try:
            # 特征标准化
            # 检查scaler是否已经fit过
            if hasattr(self.scaler, 'mean_') and self.scaler.mean_ is not None:
                # Scaler已经fit过，正常使用
                features_scaled = self.scaler.transform(features.reshape(1, -1))
            else:
                # Scaler未fit，直接使用原始特征（假设模型训练时也未使用标准化）
                features_scaled = features.reshape(1, -1)
                print("⚠️  使用未标准化的特征进行预测")
            
            # 预测类别
            predicted_class_raw = self.model.predict(features_scaled)[0]
            
            # 将数字索引转换为情感标签
            if isinstance(predicted_class_raw, (int, np.integer)):
                # 如果是数字索引，映射到情感标签
                if 0 <= predicted_class_raw < len(self.emotion_labels):
                    predicted_class = self.emotion_labels[predicted_class_raw]
                else:
                    predicted_class = f"unknown_{predicted_class_raw}"
                    print(f"⚠️  预测得到未知的情感索引: {predicted_class_raw}")
            else:
                # 如果已经是字符串标签，直接使用
                predicted_class = str(predicted_class_raw)
            
            # 获取预测概率
            if hasattr(self.model, 'predict_proba'):
                probabilities = self.model.predict_proba(features_scaled)[0]
                confidence = float(np.max(probabilities))
                
                # 构建所有类别的概率字典
                prob_dict = {}
                for i, label in enumerate(self.emotion_labels):
                    if i < len(probabilities):
                        prob_dict[label] = float(probabilities[i])
                    else:
                        prob_dict[label] = 0.0
            else:
                # KNN等模型可能没有predict_proba方法
                confidence = 1.0
                prob_dict = {label: 0.0 for label in self.emotion_labels}
                prob_dict[predicted_class] = 1.0
            
            return {
                "predicted_class": predicted_class,  # 现在确保是情感标签而不是数字
                "confidence": float(confidence),  # 确保是 Python float
                "probabilities": prob_dict  # 已经转换过了
            }
            
        except Exception as e:
            raise RuntimeError(f"预测失败: {e}")
    
    def save_model(self, model_path: str = None, scaler_path: str = None):
        """
        保存训练好的模型和标准化器
        参数：
            model_path - 模型保存路径
            scaler_path - 标准化器保存路径
        """
        if not self.is_trained:
            raise ValueError("模型尚未训练，无法保存")
        
        # 使用默认路径
        if model_path is None:
            model_path = settings.MODEL_PATH
        if scaler_path is None:
            scaler_path = settings.SCALER_PATH
        
        # 创建目录
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        os.makedirs(os.path.dirname(scaler_path), exist_ok=True)
        
        # 保存模型和标准化器
        joblib.dump(self.model, model_path)
        joblib.dump(self.scaler, scaler_path)
        
        print(f"模型已保存到: {model_path}")
        print(f"标准化器已保存到: {scaler_path}")
    
    def load_model(self, model_path: str = None, scaler_path: str = None):
        """
        加载预训练的模型和标准化器
        参数：
            model_path - 模型文件路径
            scaler_path - 标准化器文件路径 (可选)
        """
        # 使用默认路径
        if model_path is None:
            model_path = settings.MODEL_PATH
        if scaler_path is None:
            scaler_path = settings.SCALER_PATH
        
        # 检查模型文件是否存在
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"模型文件不存在: {model_path}")
        
        try:
            # 尝试加载模型文件
            model_data = joblib.load(model_path)
            
            # 检查模型文件格式
            if isinstance(model_data, dict):
                # 新格式：包含所有信息的字典
                if 'model' not in model_data:
                    raise ValueError("模型文件格式错误：缺少'model'键")
                
                self.model = model_data['model']
                
                # 如果字典中包含scaler，使用它；否则创建默认scaler
                if 'scaler' in model_data:
                    self.scaler = model_data['scaler']
                else:
                    # 创建默认的StandardScaler（适用于已训练但未保存scaler的模型）
                    from sklearn.preprocessing import StandardScaler
                    self.scaler = StandardScaler()
                    # 注意：这个scaler没有经过fit，可能需要在使用时处理
                
                # 更新模型类型和情感标签
                if 'model_type' in model_data:
                    self.model_type = model_data['model_type']
                if 'emotions' in model_data:
                    self.emotion_labels = model_data['emotions']
                
                self.is_trained = True
                print(f"✅ 从字典格式加载模型成功: {model_path}")
                print(f"   模型类型: {model_data.get('model_type', 'unknown')}")
                print(f"   支持情感: {model_data.get('emotions_count', 'unknown')}种")
                
            else:
                # 旧格式：单独的模型文件
                self.model = model_data
                
                # 尝试加载对应的scaler文件
                if os.path.exists(scaler_path):
                    self.scaler = joblib.load(scaler_path)
                    print(f"✅ 标准化器已从 {scaler_path} 加载")
                else:
                    # 创建默认scaler
                    from sklearn.preprocessing import StandardScaler
                    self.scaler = StandardScaler()
                    print(f"⚠️  标准化器文件不存在，使用默认scaler")
                
                self.is_trained = True
                print(f"✅ 从单一格式加载模型成功: {model_path}")
                
        except Exception as e:
            raise RuntimeError(f"模型加载失败: {e}")
    
    def get_feature_importance(self) -> Dict[str, float]:
        """
        获取特征重要性（仅支持RandomForest模型）
        返回值：特征重要性字典
        """
        if not self.is_trained:
            raise ValueError("模型尚未训练")
        
        if self.model_type == "random_forest":
            importance = self.model.feature_importances_
            return {"feature_importance": importance.tolist()}
        else:
            return {"message": f"{self.model_type}模型不支持特征重要性分析"}
    
    def get_model_info(self) -> Dict[str, Any]:
        """
        获取模型信息
        返回：模型类型、支持的情感等信息
        """
        if not self.is_trained:
            return {
                "type": "unknown",
                "status": "not_trained",
                "emotions": []
            }
        
        return {
            "type": self.model_type,
            "status": "trained",
            "emotions": self.emotion_labels,
            "scaler_fitted": hasattr(self.scaler, 'mean_') and self.scaler.mean_ is not None
        } 